# -*- coding: utf-8 -*-
import logging
import random

import torch
from torch.utils.data import DataLoader

# from loader import XFileLoader
from model import XModel

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class XEvaluate():
    def __init__(self, data: [DataLoader] = None, config = None):
        self.evalData = data
        self.config = config

    def eval(self, model: XModel):
        model.eval()
        statics = {"correct": 0, "wrong": 0}
        for index, test_data in enumerate(self.evalData):
            x_batch, y_batch = test_data[0]
            y_batch = y_batch.view(-1)  # 转成向量
            if torch.cuda.is_available():
                x_batch = x_batch.cuda()
                y_batch = y_batch.cuda()
            with torch.no_grad():
                y_batch_model = model.forward(x_batch)
            for y_model, y in zip(y_batch_model, y_batch):
                y_p = torch.argmax(y_model)
                if int(y_p) == int(y):
                    statics["correct"] += 1
                else:
                    statics["wrong"] += 1

        correct = statics["correct"]
        wrong = statics["wrong"]
        logger.info("预测集合条目总量：%d" % (correct + wrong))
        logger.info("预测正确条目：%d，预测错误条目：%d" % (correct, wrong))
        logger.info("预测准确率：%f" % (correct / (correct + wrong)))
        logger.info("--------------------")
